{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 6-5使用TPU训练模型\n",
    "\n",
    "如果想尝试使用Google Colab上的TPU来训练模型，也是非常方便，仅需添加6行代码。\n",
    "\n",
    "在Colab笔记本中：修改->笔记本设置->硬件加速器 中选择 TPU\n",
    "\n",
    "注：以下代码只能在Colab 上才能正确执行。\n",
    "\n",
    "可通过以下colab链接测试效果《tf_TPU》：\n",
    "\n",
    "https://colab.research.google.com/drive/1XCIhATyE1R7lq6uwFlYlRsUr5d9_-r1s\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%tensorflow_version 2.x\n",
    "import tensorflow as tf\n",
    "print(tf.__version__)\n",
    "from tensorflow.keras import * "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 一，准备数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_LEN = 300\n",
    "BATCH_SIZE = 32\n",
    "(x_train,y_train),(x_test,y_test) = datasets.reuters.load_data()\n",
    "x_train = preprocessing.sequence.pad_sequences(x_train,maxlen=MAX_LEN)\n",
    "x_test = preprocessing.sequence.pad_sequences(x_test,maxlen=MAX_LEN)\n",
    "\n",
    "MAX_WORDS = x_train.max()+1\n",
    "CAT_NUM = y_train.max()+1\n",
    "\n",
    "ds_train = tf.data.Dataset.from_tensor_slices((x_train,y_train)) \\\n",
    "          .shuffle(buffer_size = 1000).batch(BATCH_SIZE) \\\n",
    "          .prefetch(tf.data.experimental.AUTOTUNE).cache()\n",
    "   \n",
    "ds_test = tf.data.Dataset.from_tensor_slices((x_test,y_test)) \\\n",
    "          .shuffle(buffer_size = 1000).batch(BATCH_SIZE) \\\n",
    "          .prefetch(tf.data.experimental.AUTOTUNE).cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 二，定义模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.keras.backend.clear_session()\n",
    "def create_model():\n",
    "    \n",
    "    model = models.Sequential()\n",
    "\n",
    "    model.add(layers.Embedding(MAX_WORDS,7,input_length=MAX_LEN))\n",
    "    model.add(layers.Conv1D(filters = 64,kernel_size = 5,activation = \"relu\"))\n",
    "    model.add(layers.MaxPool1D(2))\n",
    "    model.add(layers.Conv1D(filters = 32,kernel_size = 3,activation = \"relu\"))\n",
    "    model.add(layers.MaxPool1D(2))\n",
    "    model.add(layers.Flatten())\n",
    "    model.add(layers.Dense(CAT_NUM,activation = \"softmax\"))\n",
    "    return(model)\n",
    "\n",
    "def compile_model(model):\n",
    "    model.compile(optimizer=optimizers.Nadam(),\n",
    "                loss=losses.SparseCategoricalCrossentropy(from_logits=True),\n",
    "                metrics=[metrics.SparseCategoricalAccuracy(),metrics.SparseTopKCategoricalAccuracy(5)]) \n",
    "    return(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 三，训练模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#增加以下6行代码\n",
    "import os\n",
    "resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])\n",
    "tf.config.experimental_connect_to_cluster(resolver)\n",
    "tf.tpu.experimental.initialize_tpu_system(resolver)\n",
    "strategy = tf.distribute.experimental.TPUStrategy(resolver)\n",
    "with strategy.scope():\n",
    "    model = create_model()\n",
    "    model.summary()\n",
    "    model = compile_model(model)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "WARNING:tensorflow:TPU system 10.26.134.242:8470 has already been initialized. Reinitializing the TPU can cause previously created variables on TPU to be lost.\n",
    "WARNING:tensorflow:TPU system 10.26.134.242:8470 has already been initialized. Reinitializing the TPU can cause previously created variables on TPU to be lost.\n",
    "INFO:tensorflow:Initializing the TPU system: 10.26.134.242:8470\n",
    "INFO:tensorflow:Initializing the TPU system: 10.26.134.242:8470\n",
    "INFO:tensorflow:Clearing out eager caches\n",
    "INFO:tensorflow:Clearing out eager caches\n",
    "INFO:tensorflow:Finished initializing TPU system.\n",
    "INFO:tensorflow:Finished initializing TPU system.\n",
    "INFO:tensorflow:Found TPU system:\n",
    "INFO:tensorflow:Found TPU system:\n",
    "INFO:tensorflow:*** Num TPU Cores: 8\n",
    "INFO:tensorflow:*** Num TPU Cores: 8\n",
    "INFO:tensorflow:*** Num TPU Workers: 1\n",
    "INFO:tensorflow:*** Num TPU Workers: 1\n",
    "INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n",
    "INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n",
    "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n",
    "Model: \"sequential\"\n",
    "_________________________________________________________________\n",
    "Layer (type)                 Output Shape              Param #   \n",
    "=================================================================\n",
    "embedding (Embedding)        (None, 300, 7)            216874    \n",
    "_________________________________________________________________\n",
    "conv1d (Conv1D)              (None, 296, 64)           2304      \n",
    "_________________________________________________________________\n",
    "max_pooling1d (MaxPooling1D) (None, 148, 64)           0         \n",
    "_________________________________________________________________\n",
    "conv1d_1 (Conv1D)            (None, 146, 32)           6176      \n",
    "_________________________________________________________________\n",
    "max_pooling1d_1 (MaxPooling1 (None, 73, 32)            0         \n",
    "_________________________________________________________________\n",
    "flatten (Flatten)            (None, 2336)              0         \n",
    "_________________________________________________________________\n",
    "dense (Dense)                (None, 46)                107502    \n",
    "=================================================================\n",
    "Total params: 332,856\n",
    "Trainable params: 332,856\n",
    "Non-trainable params: 0\n",
    "_________________________________________________________________\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = model.fit(ds_train,validation_data = ds_test,epochs = 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "Train for 281 steps, validate for 71 steps\n",
    "Epoch 1/10\n",
    "281/281 [==============================] - 12s 43ms/step - loss: 3.4466 - sparse_categorical_accuracy: 0.4332 - sparse_top_k_categorical_accuracy: 0.7180 - val_loss: 3.3179 - val_sparse_categorical_accuracy: 0.5352 - val_sparse_top_k_categorical_accuracy: 0.7195\n",
    "Epoch 2/10\n",
    "281/281 [==============================] - 6s 20ms/step - loss: 3.3251 - sparse_categorical_accuracy: 0.5405 - sparse_top_k_categorical_accuracy: 0.7302 - val_loss: 3.3082 - val_sparse_categorical_accuracy: 0.5463 - val_sparse_top_k_categorical_accuracy: 0.7235\n",
    "Epoch 3/10\n",
    "281/281 [==============================] - 6s 20ms/step - loss: 3.2961 - sparse_categorical_accuracy: 0.5729 - sparse_top_k_categorical_accuracy: 0.7280 - val_loss: 3.3026 - val_sparse_categorical_accuracy: 0.5499 - val_sparse_top_k_categorical_accuracy: 0.7217\n",
    "Epoch 4/10\n",
    "281/281 [==============================] - 5s 19ms/step - loss: 3.2751 - sparse_categorical_accuracy: 0.5924 - sparse_top_k_categorical_accuracy: 0.7276 - val_loss: 3.2957 - val_sparse_categorical_accuracy: 0.5543 - val_sparse_top_k_categorical_accuracy: 0.7217\n",
    "Epoch 5/10\n",
    "281/281 [==============================] - 5s 19ms/step - loss: 3.2655 - sparse_categorical_accuracy: 0.6008 - sparse_top_k_categorical_accuracy: 0.7290 - val_loss: 3.3022 - val_sparse_categorical_accuracy: 0.5490 - val_sparse_top_k_categorical_accuracy: 0.7231\n",
    "Epoch 6/10\n",
    "281/281 [==============================] - 5s 19ms/step - loss: 3.2616 - sparse_categorical_accuracy: 0.6041 - sparse_top_k_categorical_accuracy: 0.7295 - val_loss: 3.3015 - val_sparse_categorical_accuracy: 0.5503 - val_sparse_top_k_categorical_accuracy: 0.7235\n",
    "Epoch 7/10\n",
    "281/281 [==============================] - 6s 21ms/step - loss: 3.2595 - sparse_categorical_accuracy: 0.6059 - sparse_top_k_categorical_accuracy: 0.7322 - val_loss: 3.3064 - val_sparse_categorical_accuracy: 0.5454 - val_sparse_top_k_categorical_accuracy: 0.7266\n",
    "Epoch 8/10\n",
    "281/281 [==============================] - 6s 21ms/step - loss: 3.2591 - sparse_categorical_accuracy: 0.6063 - sparse_top_k_categorical_accuracy: 0.7327 - val_loss: 3.3025 - val_sparse_categorical_accuracy: 0.5481 - val_sparse_top_k_categorical_accuracy: 0.7231\n",
    "Epoch 9/10\n",
    "281/281 [==============================] - 5s 19ms/step - loss: 3.2588 - sparse_categorical_accuracy: 0.6062 - sparse_top_k_categorical_accuracy: 0.7332 - val_loss: 3.2992 - val_sparse_categorical_accuracy: 0.5521 - val_sparse_top_k_categorical_accuracy: 0.7257\n",
    "Epoch 10/10\n",
    "281/281 [==============================] - 5s 18ms/step - loss: 3.2577 - sparse_categorical_accuracy: 0.6073 - sparse_top_k_categorical_accuracy: 0.7363 - val_loss: 3.2981 - val_sparse_categorical_accuracy: 0.5516 - val_sparse_top_k_categorical_accuracy: 0.7306\n",
    "CPU times: user 18.9 s, sys: 3.86 s, total: 22.7 s\n",
    "Wall time: 1min 1s\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "notebook_metadata_filter": "-all",
   "text_representation": {
    "extension": ".md",
    "format_name": "markdown"
   }
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
